"""Lambda 表达式代码生成器"""

from ..core.context import CodeGenContext
from ..core.emitter import EmitResult, NodeTypeEmitter
from ...ir.expressions.lambda_expr import IRLambda
from mcpy.type_system import TypeKind, get_cpp_type
from mcpy.codegen.core.generator import get_global_generator
from .closure_generator import ClosureGenerator


class LambdaEmitter(NodeTypeEmitter):
    """Lambda 表达式代码生成器

    支持两种生成模式：
    1. Extension 模式：生成继承自 lambda_extension_base 的结构体（可存储）
    2. C++ Lambda 模式：生成原生 C++ lambda（立即调用，零开销）
    """

    def __init__(self):
        super().__init__([IRLambda])

    def emit(self, node: IRLambda, context: CodeGenContext) -> EmitResult:
        """生成 lambda 代码

        Args:
            node: Lambda IR 节点
            context: 代码生成上下文

        Returns:
            EmitResult: 生成的代码和相关信息
        """
        # 根据使用模式选择生成方式
        if node.usage_mode == "cpp_lambda":
            # 模式 1：零开销 C++ lambda
            return self._emit_cpp_lambda(node, context)
        else:
            # 模式 2：Extension lambda
            return self._emit_extension_lambda(node, context)

    def _emit_cpp_lambda(self, node: IRLambda, context: CodeGenContext) -> EmitResult:
        """生成零开销 C++ lambda

        生成示例：
            [](auto x) { return x + 1; }  // async then回调
            [](const mc::variant& x) { return x + 1; }  // 普通lambda

        Args:
            node: Lambda IR 节点
            context: 代码生成上下文

        Returns:
            EmitResult: 生成的代码
        """

        generator = get_global_generator()

        # 检查是否是async then回调
        is_async_callback = node.get_extension_param("async_then_callback", False)
        
        # 参数列表：async回调使用auto，其他使用mc::variant
        if is_async_callback:
            # 使用auto让编译器推导类型（更简洁，类型安全）
            params = ", ".join([f"auto {p.name}" for p in node.params])
        else:
            # 普通lambda使用mc::variant
            params = ", ".join([f"const mc::variant& {p.name}" for p in node.params])

        # 生成 lambda 体（使用参数名）
        lambda_context = context.create_child_context()
        lambda_context.set_extension_param(
            "in_operator_call", True, replace_existing=True
        )

        # 检查是否有await之间的语句
        between_stmts = node.get_extension_param("between_stmts", [])
        
        # 生成lambda体代码
        if between_stmts:
            # 多语句lambda：生成between_stmts + return
            statements_code = []
            for stmt in between_stmts:
                stmt_code = generator.generate_statement(stmt, lambda_context)
                statements_code.append(stmt_code)
            
            # 生成return语句
            assert node.body is not None, "Lambda body 不能为空"
            body_code = generator.generate_expression(node.body, lambda_context)
            statements_code.append(f"return {body_code};")
            
            # 组合成多行lambda体
            body_lines = "\n    ".join(statements_code)
            lambda_body = f" {{\n    {body_lines}\n}}"
        else:
            # 单语句lambda：直接return
            assert node.body is not None, "Lambda body 不能为空"
            body_code = generator.generate_expression(node.body, lambda_context)
            lambda_body = f" {{ return {body_code}; }}"

        # 生成 C++ lambda
        # async then回调：检查是否需要捕获this
        if is_async_callback:
            # 分析lambda body是否访问了成员变量
            needs_this = self._lambda_needs_this_capture(node)
            if needs_this:
                code = f"[this]({params}){lambda_body}"
            else:
                code = f"[]({params}){lambda_body}"
        else:
            code = f"[]({params}){lambda_body}"

        return EmitResult(code=code, precedence=0)

    def _lambda_needs_this_capture(self, node: IRLambda) -> bool:
        """检查lambda body是否需要捕获this
        
        如果lambda body中访问了外层函数的捕获变量，则需要this。
        
        Args:
            node: Lambda节点
            
        Returns:
            bool: 是否需要捕获this
        """
        from ...ir.expressions import IRName
        
        def visit_expr(expr):
            """递归检查表达式是否引用了捕获变量"""
            if isinstance(expr, IRName):
                # 查找符号
                symbol = expr.lookup_symbol(expr.id)
                if symbol and symbol.ir_node:
                    from ...ir.declarations import IRFunctionNode
                    # 检查是否是嵌套函数（被捕获的）
                    if isinstance(symbol.ir_node, IRFunctionNode):
                        # 查找lambda的外层函数（async函数）
                        outer_func = node.parent
                        while outer_func and not isinstance(outer_func, IRFunctionNode):
                            outer_func = outer_func.parent
                        
                        if outer_func and hasattr(outer_func, 'captured_vars'):
                            # 检查是否在捕获列表中
                            if expr.id in outer_func.captured_vars:
                                return True
            
            # 递归检查子表达式
            for child in expr.get_child_nodes():
                if visit_expr(child):
                    return True
            
            return False
        
        if node.body:
            return visit_expr(node.body)
        
        return False

    def _emit_extension_lambda(
        self, node: IRLambda, context: CodeGenContext
    ) -> EmitResult:
        """生成 Extension lambda"""
        # 1. 生成 lambda extension 结构体定义（所有模式都需要）
        struct_code = self._generate_lambda_struct(node, context)

        # 2. 将结构体定义添加到上下文（需要在函数外部）
        self._add_lambda_struct_to_context(struct_code, context, node)

        # 3. 添加必要的头文件
        context.require_include("mcpy/builtins/lambda.h")
        context.require_include("mc/variant.h")
        context.require_include("mc/memory.h")

        # 4. 检查存储模式
        storage_mode = node.get_extension_param("storage_mode")

        if storage_mode == "stack":
            # 栈模式：返回类型名（不生成创建代码）
            # 创建由 assignment_emitter 处理
            functor_name = self._get_functor_name(node)
            return EmitResult(code=functor_name, precedence=0)
        else:
            # 堆模式：生成创建代码（默认行为）
            creation_code = self._generate_lambda_creation(node, context)
            return EmitResult(code=creation_code, precedence=0)

    def _generate_lambda_struct(self, node: IRLambda, context: CodeGenContext) -> str:
        """生成 lambda extension 结构体（使用统一的 ClosureGenerator）

        Args:
            node: Lambda IR 节点
            context: 代码生成上下文

        Returns:
            str: 结构体代码
        """
        # 委托给统一的 ClosureGenerator
        generator = ClosureGenerator(
            node=node,
            context=context,
            body_fn=self._lambda_body_generator,
            name_fn=lambda n: self._get_functor_name(n),
        )
        return generator.generate()

    def _generate_lambda_creation(self, node: IRLambda, context: CodeGenContext) -> str:
        """生成 lambda 创建代码

        Args:
            node: Lambda IR 节点
            context: 代码生成上下文

        Returns:
            str: 创建代码
        """
        functor_name = self._get_functor_name(node)

        if node.captured_vars:
            # 有捕获变量，传入捕获的变量
            args = ", ".join(node.captured_vars)
            return f"mc::variant(mc::make_shared<{functor_name}>({args}))"
        else:
            # 无捕获变量
            return f"mc::variant(mc::make_shared<{functor_name}>())"

    def _get_functor_name(self, node: IRLambda) -> str:
        """获取 functor 名称（只读取，不生成）

        Args:
            node: Lambda IR 节点

        Returns:
            str: Functor 名称
        """
        # 优先使用 node.functor_name（由 type_replacement pass 设置）
        if node.functor_name:
            return node.functor_name

        # 从 extension_param 获取（由 type_replacement pass 设置）
        functor_name = node.get_extension_param("functor_name")
        if functor_name:
            node.functor_name = functor_name
            return functor_name

        # 如果走到这里，说明 type_replacement pass 没有正确执行
        # 这是错误情况，但为了调试友好，生成一个带警告的临时名称
        col_offset = node.get_extension_param("col_offset") or 0
        if hasattr(node, "source_line") and node.source_line > 0:
            temp_name = f"__lambda_{node.source_line}_{col_offset}_TEMP"
        else:
            temp_name = f"__lambda_TEMP_{id(node) & 0xFFFF}"

        # 打印警告
        print(
            f"WARNING: Lambda functor_name 未在 type_replacement pass 中设置，使用临时名称: {temp_name}"
        )

        return temp_name

    def _add_lambda_struct_to_context(
        self, struct_code: str, context: CodeGenContext, node: IRLambda
    ) -> None:
        """将 lambda 结构体定义添加到上下文

        Lambda 结构体需要定义在函数外部，因此需要添加到特殊的上下文中。

        Args:
            struct_code: 结构体代码
            context: 代码生成上下文
            node: Lambda IR 节点
        """
        # 获取 IR 模块
        ir_module = context.ir_module
        if not ir_module:
            # 尝试从节点查找
            ir_module = node.get_ir_module()

        if ir_module:
            # 将 lambda 结构体添加到模块的扩展参数中
            lambda_structs = ir_module.get_extension_param("lambda_structs") or []
            lambda_structs.append(struct_code)
            ir_module.set_extension_param("lambda_structs", lambda_structs)

    def get_priority(self) -> int:
        """Lambda 发射器的优先级"""
        return 50  # 中等优先级

    @staticmethod
    def _lambda_body_generator(node: IRLambda, context: CodeGenContext) -> str:
        """Lambda 函数体生成策略：return {expression};

        这是 Lambda 与嵌套函数的唯一差异点
        """
        from mcpy.codegen.core.generator import get_global_generator

        generator = get_global_generator()

        lambda_context = context.create_child_context()
        lambda_context.set_extension_param(
            "in_operator_call", True, replace_existing=True
        )

        assert node.body is not None, "Lambda body 不能为空"
        body_code = generator.generate_expression(node.body, lambda_context)

        return f"return {body_code};"
